#'srcnn3.onnx'文件中的Resize 是 ONNX 原生节点。其插值方式之一 bicubic 并不被 TensorRT 支持
# 为解决上述问题，我们需要创建一个新的节点替换原生 Resize 节点，并且实现新节点对应的插件代码, 插件代码指的是？
# 新改节点名称就叫 Test::DynamicTRTResize，这是种类 C++ 的写法，Test 为域名，
# 主要用于区分不同来源下的同名的节点，比如 ONNX:: 和 Test::。当然了，ONNX 本身也不存在 DynamicTRTResize 的节点名。

import torch 
from torch import nn 
from torch.nn.functional import interpolate 
import torch.onnx 
import cv2 
import numpy as np 
import os, requests

# Download checkpoint and test image 
urls = ['https://download.openmmlab.com/mmediting/restorers/srcnn/srcnn_x4k915_1x16_1000k_div2k_20200608-4186f232.pth', 
    'https://raw.githubusercontent.com/open-mmlab/mmediting/master/tests/data/face/000001.png']

names = ['srcnn.pth', 'face.png'] 
for url, name in zip(urls, names): 
    if not os.path.exists(name): 
        open(name, 'wb').write(requests.get(url).content) 

# @todo DynamicTRTResize 继承自 torch.autograd.Function，用于自定义 PyTorch 的运算符（支持前向传播和 ONNX 导出）
class DynamicTRTResize(torch.autograd.Function): 
  def __init__(self) -> None: 
        super().__init__() 
  @staticmethod 
  def symbolic(g, input, size_tensor, align_corners = False): # symbolic 方法用于定义当模型导出为 ONNX 时，此自定义操作如何转换为 ONNX 节点
      """Symbolic function for creating onnx op.""" 
      # 使用 ONNX 图构建器 g 定义了一个自定义操作符 Test::DynamicTRTResize
      return g.op( 
          'Test::DynamicTRTResize', 
          input, 
          size_tensor, 
          align_corners_i=align_corners)
  @staticmethod 
  def forward(g, input, size_tensor, align_corners = False): 
      """Run forward.""" 
      size = [size_tensor.size(-2), size_tensor.size(-1)] 
      return interpolate( 
          input, size=size, mode='bicubic', align_corners=align_corners)
class StrangeSuperResolutionNet(nn.Module): 
  def __init__(self): 
      super().__init__() 
      self.conv1 = nn.Conv2d(3, 64, kernel_size=9, padding=4) 
      self.conv2 = nn.Conv2d(64, 32, kernel_size=1, padding=0) 
      self.conv3 = nn.Conv2d(32, 3, kernel_size=5, padding=2) 
      self.relu = nn.ReLU()
  def forward(self, x, size_tensor): 
        # 使用自定义的 DynamicTRTResize 调整输入张量 x 的大小
        x = DynamicTRTResize.apply(x, size_tensor) 
        out = self.relu(self.conv1(x)) 
        out = self.relu(self.conv2(out)) 
        out = self.conv3(out) 
        return out

def init_torch_model(): 
    torch_model = StrangeSuperResolutionNet() 
    state_dict = torch.load('srcnn.pth')['state_dict'] 
    # Adapt the checkpoint：加载下载的权重文件
    for old_key in list(state_dict.keys()): 
        new_key = '.'.join(old_key.split('.')[1:]) 
        state_dict[new_key] = state_dict.pop(old_key) 
    torch_model.load_state_dict(state_dict) 
    torch_model.eval() 
    return torch_model
model = init_torch_model() 
factor = torch.rand([1, 1, 512, 512], dtype=torch.float) 
input_img = cv2.imread('face.png').astype(np.float32)

# HWC to NCHW 
input_img = np.transpose(input_img, [2, 0, 1]) 
input_img = np.expand_dims(input_img, 0)

# Inference 
torch_output = model(torch.from_numpy(input_img), factor).detach().numpy() 
# NCHW to HWC 
torch_output = np.squeeze(torch_output, 0) 
torch_output = np.clip(torch_output, 0, 255)
torch_output = np.transpose(torch_output, [1, 2, 0]).astype(np.uint8) 
# Show image 
cv2.imwrite("face_torch.png", torch_output)
x = torch.randn(1, 3, 256, 256) 
dynamic_axes={ 
        'input': { 
            0: 'batch', 
            2: 'height', 
            3: 'width' 
        }, 
        'factor': { 
            0: 'batch1', 
            2: 'height1', 
            3: 'width1' 
        }, 
        'output': { 
            0: 'batch2', 
            2: 'height2', 
            3: 'width2' 
        }, 
    }
with torch.no_grad(): 
    torch.onnx.export( 
        model, (x, factor), 
        "srcnn3.onnx", 
        opset_version=11, 
        input_names=['input', 'factor'], 
        output_names=['output'], 
        dynamic_axes=dynamic_axes)
